# -*- coding: UTF-8 -*-
# 为方便测试，请统一使用 numpy、pandas、sklearn 三种包，如果实在有特殊需求，请单独跟助教沟通
import numpy as np
import pandas as pd
from sklearn import pipeline
from sklearn.impute import SimpleImputer
from sklearn.metrics import f1_score
from sklearn import svm
import argparse
np.set_printoptions(suppress=True)
# 设定随机数种子，保证代码结果可复现
np.random.seed(1024)


class Model:
    """
    要求：
        1. 需要有__init__、train、predict三个方法，且方法的参数应与此样例相同
        2. 需要有self.X_train、self.y_train、self.X_test三个实例变量，请注意大小写
        3. 如果划分出验证集，请将实例变量命名为self.X_valid、self.y_valid
    """
    # 模型初始化，数据预处理，仅为示例
    def __init__(self, train_path, test_path):
        df_train = pd.read_csv(train_path, encoding='gbk', index_col='id')
        df_test = pd.read_csv(test_path, encoding='gbk', index_col='id')
        self.y_train = df_train['label'].values

        #RBP4列缺失值实在是太多了，删除
        df_train = df_train.drop(labels = ['RBP4','label'],axis = 1)
        df_test = df_test.drop(labels = ['RBP4', 'label'], axis = 1)

        #处理缺失值
        df_train = df_train.dropna(how = 'all')
        df_test = df_test.dropna(how = 'all')
        #离散值用众数代替缺失值
        for i in range(23):
            df_train.iloc[:,i].fillna(df_train.iloc[:,i].mode()[0],inplace = True)
            df_test.iloc[:,i].fillna(df_test.iloc[:, i].mode()[0], inplace=True)
        # print(df_train.iloc[:,81].name)
        for i in range(48,82):
            df_train.iloc[:, i].fillna(df_train.iloc[:, i].mode()[0], inplace=True)
            df_test.iloc[:, i].fillna(df_test.iloc[:, i].mode()[0], inplace=True)
        df_train.loc[:,'孕次'].fillna(df_train.loc[:,'孕次'].mode()[0],inplace = True)
        df_train.loc[:, '产次'].fillna(df_train.loc[:, '产次'].mode()[0], inplace=True)
        df_test.loc[:,'孕次'].fillna(df_test.loc[:,'孕次'].mode()[0],inplace = True)
        df_test.loc[:, '产次'].fillna(df_test.loc[:, '产次'].mode()[0], inplace=True)

        cols = []
        #将未处理的列的缺失值用平均值填充
        data_preprocessing = SimpleImputer(strategy="mean")


        self.X_train = data_preprocessing.fit_transform(df_train)
        self.X_test = data_preprocessing.transform(df_test)
        #对部分列归一化
        # print(self.X_train[:,47])
        l_age = self.X_train[:,23]
        self.X_train[:,23] =( l_age - np.mean(l_age) ) / np.std(l_age)
        l_age_2 = self.X_test[:,23]
        self.X_test[:,23] =( l_age_2 - np.mean(l_age_2) ) / np.std(l_age_2)

        for i in range(26,48):
            l_tmp = self.X_train[:,i]
            l_tmp_2 = self.X_test[:,i]
            self.X_train[:, i] = (l_tmp - np.mean(l_tmp)) / np.std(l_tmp)
            self.X_test[:,i] = (l_tmp_2 - np.mean(l_tmp_2)) / np.std(l_tmp_2)
        # np.savetxt('./test.csv', self.X_train, delimiter=',',fmt='%.04f')
        print(self.X_train)

        # https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LogisticRegression.html
        self.classification_model = svm.SVC(C = 2, kernel='rbf', gamma='auto', decision_function_shape='ovo')
        self.df_predict = pd.DataFrame(index=df_test.index)

    # 模型训练，输出训练集f1_score
    def train(self):
        self.classification_model.fit(self.X_train, self.y_train)
        y_train_pred = self.classification_model.predict(self.X_train)
        return f1_score(self.y_train, y_train_pred)

    # 模型测试，输出测试集预测结果，要求此结果为DataFrame格式，可以通过to_csv方法保存为Kaggle的提交文件
    def predict(self):
        y_test_pred = self.classification_model.predict(self.X_test)
        self.df_predict['Predicted'] = y_test_pred
        return self.df_predict


# 以下部分请勿改动！
if __name__ == '__main__':
    # 解析输入参数。在终端执行以下语句即可运行此代码： python f_model.py --train_path "f_train.csv" --test_path "f_test.csv"
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_path", type=str, default="f_train.csv", help="path to train dataset")
    parser.add_argument("--test_path", type=str, default="f_test.csv", help="path to test dataset")
    opt = parser.parse_args()

    model = Model(opt.train_path, opt.test_path)
    print('训练集维度:{}\n测试集维度:{}'.format(model.X_train.shape, model.X_test.shape))
    f1_score_train = model.train()
    print('f1_score_train={:.6f}'.format(f1_score_train))
    f_predict = model.predict()
    f_predict.to_csv('f_predict.csv')
